#pragma once

#include <dcont_arch.h>
#include <dcont/adef.h>
#include <compl.h>

/* μmp主要参考Daan Leijen大佬的实现(https://github.com/koka-lang/libmprompt)
上下文恢复汇编代码基本复用来自大佬的实现, 但在关键路径上进行了优化以提升性能，如
1. 适当改善编译器不优化手写汇编的问题
	- 使用手动优化的内联汇编, 尽可能消除依赖寄存器和访存的控制流跳转(似乎叫indirect jump?), 提升流水线效率
	- longjmp改为宏函数形式就地展开, 省去函数调用和额外保存寄存器的开销
	- 消除不必要的分支判断(setjmp后立即longjmp产生的不必要的判断), 减少分支预测失败回退的开销
	- TODO: 考虑今后若编译器__builtin_set/longjmp能够稳定可靠地使用, 可以修改宏切换到编译器实现
		这样理论上可以获得最优的性能(最好情况0切换开销)
2. 当前(2023年) CPU存在访存延迟高的问题, 优先改善cache和访存友好性
	- 上下文保存在prompt中, 这一点受hnes大佬的x86 aco实现启发(https://github.com/hnes/libaco)
		同时支持libmprompt式的栈上保存上下文的方式
	- prompt对象保存在栈顶中, 访问prompt对象同时会将栈所在页加载到dTLB, 概率降低dTLB miss, 同时也能省一个dTLB表项.
		(现代dTLB表项仅有数千个)
3. 考虑裸机可移植性, 可在裸机上编译使用(需要ABI兼容)
	- 分离prompt与内存分配器的耦合, 手动管理栈和内存
	- 可配置是否保存浮点上下文, 因为内核编程中通常不使用浮点寄存器. */

/* 参考s_task设置的语法糖 */

/* 恢复点 */
typedef fresume_t presume_t;
typedef fresume_t preturn_t;

typedef struct prompt prompt_t;

typedef	word_t	(yieldf_t)	(void *, word_t);
typedef	word_t	(promptf_t)	(prompt_t *, word_t);
typedef	word_t	(ipromptf_t)	(word_t);
/* prompt 和 stk 集成在一起，减少dtlb miss开销 */
struct prompt {
	PROMPT_CTXBUF;					/* 本prompt恢复点, x86上为64字节 */
	presume_t		*pres;			/* 需要用于指示当前恢复点*/
	struct	prompt		*parent;		/* 父prompt指针 */
	word_t			ret_arg;		/* 存放返回值或参数 */
	union{						/* 两个函数指针共用一块空间 */
		yieldf_t		*func_ptr;		/* 用于传递回调参数 */
		ipromptf_t		*cofun_ptr;		/* 用于传递协程函数 */
	};
};



/* 需要注入的依赖, 使用DEF_PROMPT */
SAI void prompt_thinit();
SAI prompt_t *prompt_cur();
SAI void prompt_setcur(prompt_t * p);

/* 算栈指针有必要吗? */
SAI void *__prompt_stkptr(prompt_t *p, bool_t downgrowth){
	paddr_t ptr = (paddr_t)p;
	return (downgrowth) ? (void *)stksize_align(ptr - PROMPT_STK_RESERVED_SIZE)
					: (void *)stksize_align(ptr + sizeof(prompt_t) + PROMPT_STK_RESERVED_SIZE);
}
SAI void *prompt_stkptr(prompt_t *p){
	return __prompt_stkptr(p, HDEF_STK_GROWS_DOWNWARD);
}

SAI void __prompt_init(prompt_t *p){
	__memset(p, 0, sizeof(prompt_t));
}
SAI prompt_t *_prompt_init(void *stk_ptr, size_t stk_size, bool_t downgrowth){
	prompt_t *p;
	/* 栈是向下增长的 */
	p = downgrowth ? (prompt_t *)((paddr_t)stk_ptr + stk_size - sizeof(prompt_t)) : \
				(prompt_t *)stk_ptr;
	__prompt_init(p);
	return p;
}
SAI prompt_t *prompt_init(void *stk_ptr, size_t stk_size){
	/* 增加初始化配置 */
	prompt_thinit();
	return _prompt_init(stk_ptr, stk_size, HDEF_STK_GROWS_DOWNWARD);
}

/* link时获取其恢复点, 同时设置其父prompt(将返回点一并设置) */
SAI PARAM_NONULL(1) presume_t *prompt_link(prompt_t *p, preturn_t *ret_point){
	// PREFETCH(p->pres);
	PROMPT_SET_RETURN(p, ret_point);
	p->parent = prompt_cur();
	prompt_setcur(p);
	/* 即将进入到p */
	return p->pres; /* FIXME: 这里*/
}
/* unlink时获取p的返回点 */
SAI PARAM_NONULL(1) preturn_t *prompt_unlink(prompt_t *p){
	PROMPT_PREFETCH(PROMPT_GET_RETURN(p->parent));	/* 由于一定会返回p的parent, 因此提前预取 */
	prompt_t *parent = p->parent;
	prompt_setcur(parent);				/* 即将返回到parent*/
	return PROMPT_GET_RETURN(p);
}

/* prompt进入 */
#define prompt_enter(p, fp, arg)			\
do { promptf_t *pfunc = (fp); setstack_exec2(pfunc, p, arg, prompt_stkptr(p)); } while(0)

/* setexec 可能丢失浮点寄存器? 得提前保存 */

/* exec 可以弄一层保护函数，确保用户正常返回 */
/* prompt_exec
--------------------------------------------------*/
#define __prompt_exec(prompt, fp, arg, cofunc, coarg, RETURN_ROUTINE)  ({\
	prompt_t *cur = prompt_cur();					\
	preturn_t *ret;							\
	PROMPT_GET_CTXBUF(cur, ret);					\
	prompt_t *p = (prompt);						\
	/* 进入前进行连接 */						\
	prompt_link(p, ret);						\
	ipromptf_t *cofun = (cofunc);					\
	if(cofun){/*传空可以被编译器优化*/				\
		p->cofun_ptr = cofun;					\
		p->ret_arg = coarg;					\
	}								\
	if(SETJMP(ret) == 0){						\
		prompt_enter(p, fp, arg);				\
	}else{								\
		RETURN_ROUTINE(p, ret);					\
	}								\
	p->ret_arg;	/*不读可以被优化*/				\
})

#define __prompt_return_donothing(p, ret)
#define __prompt_return_custom(p, ret)					\
if (p->func_ptr) { p->ret_arg = ((prompt_func_t *)(p->func_ptr))(p, p->ret_arg);  p->func_ptr = NULL_PTR; }

#define _prompt_exec(prompt, fp, arg, cofunc, coarg)	\
	__prompt_exec(prompt, fp, arg, cofunc, coarg, __prompt_return_donothing)
#define _prompt_execf(prompt, fp, arg, cofunc, coarg)	\
	__prompt_exec(prompt, fp, arg, cofunc, coarg, __prompt_return_custom)

/* 普通执行 */
#define prompt_exec(prompt, fp, arg)	\
	_prompt_exec(prompt, fp, arg, NULL_PTR, 0)
#define prompt_execf(prompt, fp, arg)	\
	_prompt_exec(prompt, fp, arg, NULL_PTR, 0)
/* 封装执行 */
#define prompt_wrapexec(prompt, func, arg, wrapfunc, wraparg)	\
	_prompt_exec(prompt, fp, arg, wrapfunc, wraparg)
#define prompt_wrapexecf(prompt, func, arg, wrapfunc, wraparg)	\
	_prompt_execf(prompt, fp, arg, wrapfunc, wraparg)

/* prompt_return
--------------------------------------------------*/
#define __prompt_return(p, yieldf)				\
	preturn_t *retp = prompt_unlink(p);			\
	if(yieldf) { yieldf_t *yf = (yieldf);			\
	p->func_ptr = (yf); }/* 可以被优化掉*/			\
	LONGJMP(retp, 1);					


#define __prompt_return_val(prompt, yieldf, ret_val)	do{	\
	prompt_t *__p = (prompt);				\
	__p->ret_arg = (ret_val);				\
	__prompt_return(__p, yieldf);				\
}while(0)

/* 返回 */
#define prompt_return_v(prompt, ret_val)		__prompt_return_val(prompt, NULL_PTR, ret_val);
#define prompt_returnf_v(prompt, yieldf, ret_val)	__prompt_return_val(prompt, yieldf, ret_val);

#define prompt_return(prompt)				__prompt_return(prompt, NULL_PTR)
#define prompt_returnf(prompt, yieldf)			__prompt_return(prompt, yieldf)

/* prompt_yield
--------------------------------------------------*/
#define __prompt_switch(a, b, use_ctxsw)		\
	if (use_ctxsw) { __ctxsw(a, b); } else{ if(SETJMP(a) == 0){ LONGJMP(b, 1); } }
/* once 使用__ctxsw原语，省一次if分支，优化分支预测 */
#define __prompt_yieldf(pres, cur, yto, fun, arg, use_ctxsw)	\
	preturn_t *ret = prompt_unlink(yto);			\
	yieldf_t *func = (fun);					\
	if (func) yto->func_ptr = func; /* 可优化 */	\
	__prompt_switch(pres, ret, use_ctxsw)			\
	yto->ret_arg


#define _prompt_yieldf_val(p, fun, arg, use_ctxsw)	({	\
	prompt_t *cur = prompt_cur();				\
	presume_t *res;						\
	PROMPT_GET_CTXBUF(cur, res);				\
	prompt_t *yto = (p);					\
	yto->ret_arg = (arg);					\
	yto->pres = res;/* FIXME: 注意啊 */			\
	__prompt_yieldf(res, cur, yto, fun, arg, use_ctxsw);	\
})

#define _prompt_yieldf(p, fun, use_ctxsw)		({	\
	prompt_t *cur = prompt_cur();				\
	presume_t *res;						\
	PROMPT_GET_CTXBUF(cur, res);				\
	prompt_t *yto = (p);					\
	yto->pres = res;/* FIXME: 注意啊 */			\
	__prompt_yieldf(res, cur, yto, fun, arg, use_ctxsw);	\
})

/* a yield b ->a(fpu) b ->b  resume a -> */
#define prompt_yield_v(p, arg)				_prompt_yieldf_val(p, NULL_PTR, arg, 1)
#define prompt_yieldf_v(p, f, arg)			_prompt_yieldf_val(p, f, arg, 1)

#define prompt_yield(prompt)				_prompt_yieldf(prompt, NULL_PTR, 1)
#define prompt_yieldf(prompt, yieldf)			_prompt_yieldf(prompt, yieldf, 1)

/* prompt_resume
--------------------------------------------------*/
#define ___prompt_resume(pret, p, RETURN_ROUTINE, use_ctxsw)	\
	PROMPT_PREFETCH(&p->pres);				\
	presume_t *resto = prompt_link(p, pret);		\
	__prompt_switch(pret, resto, use_ctxsw);		\
	RETURN_ROUTINE(p, pret);				\
	/* 不接收变量就不会读 */				\
	p->ret_arg

#define __prompt_resume_val(prompt, arg, RETURN_ROUTINE, use_ctxsw)\
({								\
	prompt_t *p = (prompt);					\
	prompt_t *cur = prompt_cur();				\
	preturn_t *pret;					\
	PROMPT_GET_CTXBUF(cur, pret);				\
	p->ret_arg = arg;					\
	___prompt_resume(pret, p, RETURN_ROUTINE, use_ctxsw);	\
})

#define __prompt_resume(prompt,RETURN_ROUTINE, use_ctxsw) ({	\
	prompt_t *p = (prompt);					\
	prompt_t *cur = prompt_cur();				\
	preturn_t *pret;					\
	PROMPT_GET_CTXBUF(cur, pret);				\
	___prompt_resume(pret, p, RETURN_ROUTINE, use_ctxsw);	\
})


#define _prompt_resume(p, arg, use_ctxsw)		__prompt_resume_val(p, arg, __prompt_return_donothing, use_ctxsw)
#define _prompt_resumeyf(p, arg, use_ctxsw)		__prompt_resume_val(p, arg, __prompt_return_custom, use_ctxsw)

#define prompt_resume_v(p, arg)				_prompt_resume(p, arg, 1)
#define prompt_resumef_v(p, arg)			_prompt_resumeyf(p, arg, 1)

#define prompt_resume(p)				__prompt_resume(p, __prompt_return_donothing, 1)
#define prompt_resumef(p)				__prompt_resume(p, __prompt_return_custom, 1)

/* prompt wrapper
--------------------------------------------------*/
#define __PROMPTV(func)					\
static PROMPT_FUNC 					\
word_t __##func##_prompt(prompt_t *p, word_t arg){	\
	prompt_t *c = p;				\
	ipromptf_t *fun = (func);			\
	word_t ret = fun(arg);				\
	prompt_return_v(c, ret); }

#define __PROMPT(func)					\
static PROMPT_FUNC					\
word_t __##func##_prompt(prompt_t *p, word_t arg){	\
	prompt_t *c = p;				\
	ipromptf_t *fun = (func);			\
	fun(arg);					\
	prompt_return(c); }

#define DEF_PROMPT_FUNC_RETV(func)			__PROMPTV(func)
#define DEF_PROMPT_FUNC(func)				__PROMPT(func)
#define PROMPT_BYFUNC(func)				__##func##_prompt

#define DEF_PROMPT_MODULE(cur_prompt, root_prompt)	\
SAI prompt_t *prompt_cur(){ return (cur_prompt); }	\
SAI void prompt_setcur(prompt_t *p){ (cur_prompt) = p; }\
SAI prompt_t *prompt_root(){ return (root_prompt); }	\
SAI void prompt_thinit(){ if (!cur_prompt) prompt_setcur(root_prompt); }

/* enter时执行 root -> prompt_enter(p) -> 获得p的栈大小和栈指针 -> [在root栈分配current_resume点] -> link -> 设置p的ret -> setstack_exec
prompt_return(p, arg) (yield)

prompt_yield(p, arg) -> [unlinked_p 获得p的ret点] -> 将arg填入p的ret_arg -> [在本栈分配current_resume点]
-> [setjmp] -> [将p的resume点设置为current_resume] -> [跳转到p的ret点]
|
[获得p的返回值]

prompt_resume(p, arg) -> [访问p获得其res点] -> 将arg填入p的ret_arg -> [在本栈分配current_resume点] -> link 设置p的parent
-> [将p的ret点设置在current_resume点] -> [跳转到p的res点]
|
[获得p的返回值]
*/

